from transformers import BertModel

import torch

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 加载hugging face提供的模型 bert-base-chinese
pretrained = BertModel.from_pretrained(
    r"../../my_model_cache/bert-base-chinese/models--bert-base-chinese/snapshots/c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f").to(
    DEVICE)
print(pretrained)


# 定义下游任务模型（将主干网络所提取的特征分类）
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 定义一个全连接层，用于将 BERT 的输出映射到分类结果。BERT 的输出向量维度为 768，
        # 因此输入到 fc 层的维度为 768，输出维度为分类的类别数（这里是 8 分类任务）。
        self.fc = torch.nn.Linear(768, 8)  # 8分类

    def forward(self, input_ids, attention_mask, token_type_ids):
        with torch.no_grad():
            out = pretrained(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
        out = self.fc(out.last_hidden_state[:, 0])
        out = out.softmax(dim=1)
        return out

print("==========================")
mode = Model()
print(mode)